"""
Task training callback for continual learning.

This callback provides the core functionality for sequential task training,
coordinating with TaskBoundaryCallback for proper task management.
"""

from typing import Dict, List, Optional, Union, Any, Sequence
import pytorch_lightning as pl
from pytorch_lightning.callbacks import Callback
from torch.utils.data import DataLoader

from ..module.continual_module import ContinualModule
from .task_boundary_callback import TaskBoundaryCallback


class TaskTrainingCallback(Callback):
    """
    核心任务训练回调，负责编排多任务顺序训练。

    这个回调是框架的核心，它协调TaskBoundaryCallback来实现：
    1. 多任务的顺序训练
    2. 任务间的状态管理
    3. 与Lightning标准训练流程的集成

    Args:
        task_boundary_callback: 任务边界管理器，负责任务切换通知
        reset_optimizer_between_tasks: 是否在任务间重置优化器状态
    """

    def __init__(
        self,
        task_boundary_callback: Optional[TaskBoundaryCallback] = None,
        reset_optimizer_between_tasks: bool = False,
    ):
        super().__init__()
        self.task_boundary_callback = task_boundary_callback or TaskBoundaryCallback()
        self.reset_optimizer_between_tasks = reset_optimizer_between_tasks

        # 训练状态追踪
        self.completed_tasks = set()
        self.task_results = {}

    def train_multiple_tasks(
        self,
        trainer: pl.Trainer,
        model: ContinualModule,
        tasks_data: List[Dict[str, Union[DataLoader, Sequence[DataLoader]]]],
    ) -> List[Dict[str, Any]]:
        """
        顺序训练多个任务。

        这是框架的主要入口点，用户通过这个方法来执行多任务学习。

        Args:
            trainer: PyTorch Lightning trainer
            model: ContinualModule实例
            tasks_data: 任务数据列表，每个元素包含'train'和'val'数据加载器

        Returns:
            每个任务的训练结果列表
        """
        results = []

        for task_id, task_data in enumerate(tasks_data):
            print(f"\n=== 开始训练任务 {task_id} ===")

            # 设置当前任务
            self.task_boundary_callback.set_task(task_id, trainer, model)

            # 重置优化器（如果需要）
            if self.reset_optimizer_between_tasks and task_id > 0:
                self._reset_optimizer_state(trainer)

            # 执行标准Lightning训练
            train_loader = task_data["train"]
            val_loader = task_data.get("val", None)

            result = trainer.fit(
                model, train_dataloaders=train_loader, val_dataloaders=val_loader
            )

            # 记录结果
            self.completed_tasks.add(task_id)
            self.task_results[task_id] = result
            results.append(result)

            print(f"=== 任务 {task_id} 训练完成 ===\n")

        return results

    def test_all_tasks(
        self,
        trainer: pl.Trainer,
        model: ContinualModule,
        test_dataloaders: List[DataLoader],
    ) -> List[Dict[str, Any]]:
        """
        在所有任务上测试模型。

        Args:
            trainer: PyTorch Lightning trainer
            model: 训练好的模型
            test_dataloaders: 每个任务的测试数据加载器

        Returns:
            每个任务的测试结果
        """
        all_results = []

        for task_id, test_loader in enumerate(test_dataloaders):
            print(f"测试任务 {task_id}...")

            # 设置当前任务
            self.task_boundary_callback.set_task(task_id, trainer, model)

            # 执行测试
            results = trainer.test(model, dataloaders=test_loader)
            all_results.append(results[0] if results else {})

        return all_results

    def _reset_optimizer_state(self, trainer: pl.Trainer) -> None:
        """重置优化器状态，防止任务间干扰。"""
        if not trainer.optimizers:
            return

        from collections import defaultdict

        for optimizer in trainer.optimizers:
            optimizer.state = defaultdict(dict)

        print("已重置优化器状态")
